import os
import pickle
import urllib.request

import numpy as np

from core.data.dataset import Dataset
from implement.transforms.astype import ToFloat
from implement.transforms.compose import Compose
from implement.transforms.flatten import Flatten
from implement.transforms.normalize import Normalize

import gzip
import matplotlib.pyplot as plt

from utils.file_opera import get_file
from utils.os_main import get_file_path_in_project, does_directory_exist, create_directory


class PTB(Dataset):
    """
        PTB语料库经常被用作评价提案方法的基准
        在PTB语料库中，一行保存一个句子
    """
    def __init__(self, train=True,
                 transform=Compose([Flatten(), ToFloat(),
                                    Normalize(0., 255.)]),
                 target_transform=None):
        super().__init__(train, transform, target_transform)
        self.dataset_dir = None
        self.vocab_file = None
        self.save_file = None
        self.key_file = None
        self.url_base = None

        self.prepare()

    def prepare(self):
        self.url_base = 'https://raw.githubusercontent.com/tomsercu/lstm/master/data/'
        self.key_file = {
            'train': 'ptb.train.txt',
            'test': 'ptb.test.txt',
            'valid': 'ptb.valid.txt'
        }
        self.save_file = {
            'train': 'ptb.train.npy',
            'test': 'ptb.test.npy',
            'valid': 'ptb.valid.npy'
        }
        self.vocab_file = 'ptb.vocab.pkl'
        if not does_directory_exist(get_file_path_in_project('temp')):
            create_directory(get_file_path_in_project('temp'))

        self.dataset_dir = get_file_path_in_project('temp')

    def _download(self, file_name):
        """
        下载指定文件，如果文件已存在则直接返回。

        Args:
            file_name (str): 要下载的文件名。

        Returns:
            None
        """
        # 构建文件路径
        file_path = self.dataset_dir + '/' + file_name

        # 如果文件已存在，则直接返回
        if os.path.exists(file_path):
            return

        # 输出下载提示信息
        print('Downloading ' + file_name + ' ... ')

        try:
            # 尝试下载文件
            urllib.request.urlretrieve(self.url_base + file_name, file_path)
        except urllib.error.URLError:
            # 如果遇到 URLError，则尝试创建未验证的 HTTPS 上下文
            import ssl
            ssl._create_default_https_context = ssl._create_unverified_context
            urllib.request.urlretrieve(self.url_base + file_name, file_path)

        # 输出下载完成信息
        print('Done')

    def load_vocab(self):
        # 拼接词汇表路径
        vocab_path = os.path.join(self.dataset_dir, self.vocab_file)

        # 如果词汇表已存在，直接加载并返回
        if os.path.exists(vocab_path):
            with open(vocab_path, 'rb') as f:
                """
                    pickle 是 Python 中的一个模块，用于在对象之间进行序列化和反序列化。
                    序列化是将对象转换为字节流的过程，而反序列化是从字节流中恢复对象的过程。
                    在上述代码中，pickle 被用来将词汇表 (word_to_id 和 id_to_word) 序列化并保存到文件，
                    以便在下次运行时直接加载，从而避免重复构建词汇表的过程，提高程序的效率。
                """
                word_to_id, id_to_word = pickle.load(f)
            return word_to_id, id_to_word

        # 初始化词典
        word_to_id = {}
        id_to_word = {}
        data_type = 'train'
        file_name = self.key_file[data_type]
        file_path = os.path.join(self.dataset_dir, file_name)

        # 下载数据集文件
        self._download(file_name)

        # 读取文件内容，将换行符替换为'<eos>'，去除首尾空格后切分为单词列表
        words = open(file_path).read().replace('\n', '<eos>').strip().split()

        # 构建词典，记录单词到ID的映射关系
        for i, word in enumerate(words):
            if word not in word_to_id:
                tmp_id = len(word_to_id)
                word_to_id[word] = tmp_id
                id_to_word[tmp_id] = word

        # 将词典保存到文件
        with open(vocab_path, 'wb') as f:
            pickle.dump((word_to_id, id_to_word), f)

        return word_to_id, id_to_word

    def load_data(self, data_type='train'):
        """
        加载数据并返回语料库、词汇表的映射关系。

        Args:
            data_type (str): 数据的种类，可选值为 'train', 'test' 或 'valid'（val）。

        Returns:
            np.ndarray: 语料库，表示为单词 ID 的数组。
            dict: 单词到 ID 的映射。
            dict: ID 到单词的映射。
        """
        # 将 'val' 转换为 'valid'
        if data_type == 'val':
            data_type = 'valid'

        # 保存路径
        save_path = self.dataset_dir + '/' + self.save_file[data_type]

        # 加载词汇表
        word_to_id, id_to_word = self.load_vocab()

        # 如果已存在保存的语料库文件，则直接加载
        if os.path.exists(save_path):
            corpus = np.load(save_path)
            return corpus, word_to_id, id_to_word

        # 读取数据文件并进行预处理
        file_name = self.key_file[data_type]
        file_path = self.dataset_dir + '/' + file_name
        self._download(file_name)

        words = open(file_path).read().replace('\n', '<eos>').strip().split()
        corpus = np.array([word_to_id[w] for w in words])

        # 保存处理后的语料库到文件
        np.save(save_path, corpus)
        return corpus, word_to_id, id_to_word


if __name__ == '__main__':
    # 准备ptb数据
    ptb = PTB()
    ptb.prepare()
    for data_type in ('train', 'val', 'test'):
        ptb.load_data(data_type)